import robomaster
from robomaster import robot

from parameters import *

import time


class MyRobot:
    def __init__(self, manager):
        self.robot = robot.Robot()
        self.robot.initialize(conn_type='ap')

        # check if successfully connect to the robot
        version = self.robot.get_version()
        if version:
            print('Successfully connected to the robot.')
        else:
            print('Failed to connect.')
            exit()

        self.robot.set_robot_mode(mode='chassis_lead')

        self.sensor = self.robot.sensor
        self.sensor.sub_distance(freq=update_freq, callback=self.infrared_callback)

        self.chassis = self.robot.chassis

        self.pubsub_manager = manager
        
        self.infrared_publisher = self.pubsub_manager.create_publisher('infrared')

        self.movement_subscriber = self.pubsub_manager.create_subscription('movement', self.movement_callback)

        self.lock = False

    def infrared_callback(self, data):
        distance = int(data[0]) // grid_size
        distance = 1 if distance == 0 else distance
        msg = distance
        if self.lock:
            return
        self.lock = True
        self.infrared_publisher.publish(msg)
        print('node robot publishes infrared distance', msg)      

    def movement_callback(self, msg):
        print('node robot receives movement', msg)
        self.lock = True

        turning = msg
        if turning == 0:
            # move forward for a grid
            self.chassis.move(x=grid_size / 1000, y=0, z=0, xy_speed=1)
            print('robot move forward for one grid')
            time.sleep(2)
        elif turning == 1:
            # turn right
            self.chassis.move(x=0, y=0, z=-90, z_speed=45)
            print('robot turns right')
            time.sleep(3)
        elif turning == 2:
            # turn back
            self.chassis.move(x=0, y=0, z=180, z_speed=45)
            print('robot turns back')
            time.sleep(6)
        elif turning == 3:
            # turn left
            self.chassis.move(x=0, y=0, z=90, z_speed=45)
            print('robot turns left')
            time.sleep(3)
        else:
            pass
        
        self.lock = False
